package IC.asm.RegAllocation;

import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.SortedSet;
import java.util.TreeSet;

import IC.asm.ASMEntry;
import IC.asm.Instruction;
import IC.asm.Instruction.Ops;
import IC.asm.decl.*;
import IC.asm.operands.*;
import IC.asm.operands.MachineReg.Regs;


public class RegisterAllocator extends DefaultAllocator<Object,Object> {

	// list of free registers
	private List<RegOperand> freeRegs;
	
	private Map<VirtualReg, MachineReg> regAssoc;
	
	private RegisterAllocator() {
		this.freeRegs = new ArrayList<RegOperand>();
		
		this.regAssoc = new HashMap<VirtualReg, MachineReg>();
		
		// init free reg list
		freeRegs.add(new MachineReg(Regs.ebx));
		freeRegs.add(new MachineReg(Regs.ecx));
		freeRegs.add(new MachineReg(Regs.esi));
		freeRegs.add(new MachineReg(Regs.edi));
		
	}
	
	public static List<ASMEntry> Process(List<ASMEntry> asmList) {
		RegisterAllocator ra = new RegisterAllocator();
		for(ASMEntry e : asmList) {
			e.accept(ra, null);
		}
		return code;
	}
	

	private MachineReg getRegister(VirtualReg vreg) {
		MachineReg mreg;
		if(regAssoc.get(vreg) != null) {
			// virtual reg is already associated to a machine reg
			// return it
			return regAssoc.get(vreg);
		}
		else {
			assert(!freeRegs.isEmpty());
			// get machine register from free regs, remove from free list
			mreg = (MachineReg)freeRegs.get(0); freeRegs.remove(mreg);
			
			// associate from now on with virtual reg
			regAssoc.put(vreg, mreg);
			
			// return machine reg for usage
			return mreg;
		}
	}
	
	public String visit(Instruction that, Object d) {

		RegOperand reg;
		MachineReg mreg;
		
		Instruction inst = that;
		
		if(inst.getSource() instanceof VirtualReg){
			mreg = getRegister((VirtualReg)inst.getSource());
			inst = new Instruction(inst.getOp(), mreg, inst.getDest());
		}
		else if(inst.getSource() instanceof MemOperand) {
			MemOperand memOp = ((MemOperand)inst.getSource());
			RegOperand base = memOp.getBase();
			RegOperand index = memOp.getIndex();
			
			if(base	instanceof VirtualReg) {
				base = getRegister((VirtualReg)base);
			}
			if(index instanceof VirtualReg) {
				index = getRegister((VirtualReg)index);
			}
	
			memOp = new MemOperand(base, memOp.getOffset(), index, memOp.getScale());
			inst = new Instruction(inst.getOp(), memOp, inst.getDest());
		}

		if(inst.getDest() instanceof VirtualReg) {		
			mreg = getRegister((VirtualReg)inst.getDest());
			inst = new Instruction(inst.getOp(), inst.getSource(), mreg);
		}
		else if(inst.getDest() instanceof MemOperand) {
			MemOperand memOp = ((MemOperand)inst.getDest());
			RegOperand base = memOp.getBase();
			RegOperand index = memOp.getIndex();

			if(base	instanceof VirtualReg) {
				base = getRegister((VirtualReg)base);
			}
			if(index instanceof VirtualReg) {
				index = getRegister((VirtualReg)index);
			}
			
			memOp = new MemOperand(base, memOp.getOffset(), index, memOp.getScale());
			inst = new Instruction(inst.getOp(), inst.getSource(), memOp);
		}
	
		
		addCode(inst);
		
		return null;
	}

	public String visit(MachineReg that, Object d) {
		if(freeRegs.contains(that)) {
			freeRegs.remove(that);
		}
		return null;
	}

	public String visit(VirtualReg that, Object d) {
		assert(!freeRegs.isEmpty());

		return null;
	}
}
